import cv2
import numpy as np
import torchvision
from torch.utils.data import DataLoader
from torchvision.models import vgg16_bn
import torch
import torch.nn as nn
import os
from torchvision.transforms import transforms
from PIL import Image


# 加载数据集
def get_dataloader(mode):
    """
    获取数据集加载
    :param mode:
    :return:
    """
    # 准备数据迭代器
    # 这里我已经下载好了，所以是否需要下载写的是false
    # 准备数据集，其中0.1307，0.3081为MNIST数据的均值和标准差，这样操作能够对其进行标准化
    # 因为MNIST只有一个通道（黑白图片）,所以元组中只有一个值
    dataset = torchvision.datasets.MNIST('./mini',
                                         train=mode,
                                         download=False,
                                         transform=torchvision.transforms.Compose([
                                             torchvision.transforms.ToTensor(),
                                             torchvision.transforms.Normalize(
                                                 (0.1307,), (0.3081,))
                                         ]))

    return DataLoader(dataset, batch_size=64, shuffle=True)


# 图像预处理函数，将图像转换成[224,224]大小,并进行Normalize,返回[1,3.224,224]的四维张量
def preprocess_image(cv2im, resize_im=True):
    # 在ImageNet100万张图像上计算得到的图像的均值和标准差，它会使图像像素值大小在[2.7.2.1]之间，但是整体图像像素值的分布会是标准正态分布(（均值为0，方差为1)
    # 之所以使用这种方法，是因为这是基于 lmageNet的预训练VGG16对输入图像的要求
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]

    # 改变图像大小并进行Normalize
    if resize_im:
        cv2im = cv2.resize(cv2im, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
    im_as_arr = np.float32(cv2im)
    im_as_arr = np.ascontiguousarray(im_as_arr[..., ::-1])
    im_as_arr = im_as_arr.transpose(2, 0, 1)  # 将[W,H,CJ的次序改变为[C,W,H]

    for channel, _ in enumerate(im_as_arr):  # 进行在ImageNet上预训练的VGG16要求的ImageNe输入图像的Normalizeim_as_arr[channel] /= 255
        im_as_arr[channel] /= 255
        im_as_arr[channel] -= mean[channel]
        im_as_arr[channel] /= std[channel]

    # 转变为三维 Tensor,[C,W, H
    im_as_ten = torch.from_numpy(im_as_arr).float()
    im_as_ten = im_as_ten.unsqueeze_(0)  # 扩充为四维Tensor,变为[1,C,W,H

    return im_as_ten  # 返回处理好的[1,3,224,224]四维Tensor


class FeaureVisualization():
    def __init__(self, img_path, selected_layer):
        """

        :param img_path: 图片的地址
        :param selected_layer: 期望输出的特征图
        """
        self.img_path = img_path
        self.select_layer = selected_layer
        # 如果预训练模型已经下载好了，就加载本地
        if os.path.exists("./model/vgg16.pth"):
            self.model = vgg16_bn(pretrained=False)
            self.model.load_state_dict(torch.load(r"./model/vgg16.pth"))
            self.model = self.model.features
        else:
            self.model = vgg16_bn(pretrained=True).features

    def preprocess_image2(self):
        """
        加载图片，把图片进行归一化处理，变成[0,1]之间
        返回[1,C,H,W]格式的tensor形式
        :return:
        """
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        T = transforms.Compose([
            transforms.Resize(size=(224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
        return T(Image.open(self.img_path)).unsqueeze(0)

    def get_features(self):
        """
        获取特征图信息
        :return:
        """
        x = self.preprocess_image2()  # 获取图片信息
        print(x.shape)
        for index, layer in enumerate(self.model):
            x = layer(x)
            if index == self.select_layer:
                return x

    def get_single_feature(self):
        return self.get_features()

    def save_feature_to_img(self):
        features = self.get_features()
        print(features.shape)
        for i in range(features.shape[1]):

            feature = features[:, i, :, :].squeeze().data.numpy()
            # batch为1,所以所以直接view成二维的
            # #根据图像的像素值中最大最小值，将特征图的像素值归一化到了[0,1];
            feature = (feature - np.amin(feature)) / (np.amax(feature) - np.amin(feature + 1e-5))
            feature = np.round(feature * 255)

            if not os.path.exists("./image"):
                os.path.exists("./image")

            cv2.imwrite("./image/{}.jpg".format(i), feature)


if __name__ == '__main__':
    myClass = FeaureVisualization("./input.jpg", 0)
    print(myClass.model)

    myClass.save_feature_to_img()
